from utils import *
from common import *
import time
import sys

# ===================== 客户端 =====================
class ControlClient:
    def __init__(self, master_host, master_port, client_id):
        self.master_host = master_host
        self.master_port = master_port
        self.client_id = client_id

    def connect(self):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.connect((self.master_host, self.master_port))

    def send_command(self, command, **kwargs):
        msg = {'type': MessageType.ControlCommand, 'command': command}
        msg.update(kwargs)
        MessageProtocol.send(self.sock, msg)

    def add_tasks(self, tasks: list[Task]):
        self.send_command('add_tasks', tasks=[t.model_dump() for t in tasks])

    def stop_server(self):
        if self.sock: self.send_command('stop_server')
        
    def get_slave_info(self):
        self.send_command('get_slave_info')
        response = MessageProtocol.recv(self.sock)
        if not response:
            print("远程主机断开")
        
        msg_type = response['type']
        if msg_type == MessageType.SlaveInfo:
            print(f"{json.dumps(response['data'], ensure_ascii=False, indent=4)}")
        elif msg_type == MessageType.Refuse:
            print(f"{response['message']}")
        elif msg_type == MessageType.Exception:
            print(f"{response['message']}")
            
    def get_task_info(self):
        self.send_command('get_task_info')
        response = MessageProtocol.recv(self.sock)
        if not response:
            print("远程主机断开")
        
        msg_type = response['type']
        if msg_type == MessageType.TaskInfo:
            print(f"{json.dumps(response['data'], ensure_ascii=False, indent=4)}")
        elif msg_type == MessageType.Refuse:
            print(f"{response['message']}")
        elif msg_type == MessageType.Exception:
            print(f"{response['message']}")
            
    def get_task_results(self, id):
        self.send_command('get_task_results', id=id)
        response = MessageProtocol.recv(self.sock)
        if not response:
            print("远程主机断开")
        
        msg_type = response['type']
        if msg_type == MessageType.TaskResult:
            print(f"{json.dumps(response['data'], ensure_ascii=False, indent=4)}")
        elif msg_type == MessageType.Refuse:
            print(f"{response['message']}")
        elif msg_type == MessageType.Exception:
            print(f"{response['message']}")

    def close(self):
        
        self.sock.close()
        self.sock = None
        
def client_start():
    # 启动控制客户端
    control_client = ControlClient(master_host='127.0.0.1', master_port=9000, client_id='control_client_1')
    control_client.connect()
    
    while True:
        print("""请选择：
              1. 添加任务
              2. 获取slave信息
              3. 获取task信息
              4. 获取任务结果
              5. 停止服务器
              6. 退出""")
        opt = int(input("input: "))
        if opt == 1:
            # 添加任务
            control_client.add_tasks([
                Task(id="1", dataset={"bbh": ["bbh-asd", "bbh-aaa"]}, model='deepseek'),
                Task(id="1", dataset={"mmlu": ["mmlu-asd", "mmlu-aaa"]}, model='qwen'),
                Task(id="1", dataset={"story": ["story-asd"]}, model='gemini'),
            ])
            control_client.add_tasks([
                Task(id="2", dataset={"bbh": ["bbh-asd", "bbh-aaa"]}, model='deepseek'),
                Task(id="2", dataset={"mmlu": ["mmlu-asd", "mmlu-aaa"]}, model='qwen'),
                Task(id="2", dataset={"story": ["story-asd"]}, model='gemini'),
            ])
        elif opt == 2:
            control_client.get_slave_info()
        elif opt == 3:
            control_client.get_task_info()
        elif opt == 4:
            id = input("输入查询id")
            control_client.get_task_results(id)
        elif opt == 5:
            control_client.stop_server()
            return
        else:
            print('请重新输入')
    
# ==================== 使用示例 ====================
if __name__ == '__main__':
    client_start()